// Author: Ivan Chau
package uk.ac.cam.kcic2.photonMap;

import java.awt.Color;
import java.awt.image.BufferedImage;
import java.io.File;
import java.io.IOException;
import java.util.ArrayDeque;
import java.util.ArrayList;
import java.util.Deque;

import javax.imageio.ImageIO;

import uk.ac.cam.kcic2.math.Params;
import uk.ac.cam.kcic2.math.Vector3D;
import uk.ac.cam.kcic2.rayTracer.RayTracer;
import uk.ac.cam.kcic2.scene.Colour;
import uk.ac.cam.kcic2.scene.ColourCoeffs;
import uk.ac.cam.kcic2.scene.IntersectionInfo;
import uk.ac.cam.kcic2.scene.LightSource;
import uk.ac.cam.kcic2.scene.Ray;
import uk.ac.cam.kcic2.scene.Scene;

public class PhotonMapper extends RayTracer{

	private final int maxPhotonsPerLight;
	private final int maxCausticsPhotonsPerLight;
	private PhotonMap globalMap;
	private PhotonMap causticsMap;
	private boolean[][] angles = new boolean[256][256];
	private ArrayList<Integer> phiIndex = new ArrayList<Integer>();
	private ArrayList<Integer> thetaIndex = new ArrayList<Integer>();
	private double radius;

	public PhotonMapper(Scene sceneSettings, double screenDistance, 
			double width, double height, int hPixels, int vPixels, int bounceNum, int photonsPerLight, 
			int causticsNum, double r){
		super(sceneSettings, screenDistance, width, height, hPixels, vPixels, bounceNum);
		maxPhotonsPerLight = photonsPerLight;
		maxCausticsPhotonsPerLight = causticsNum;
		radius = r;

		globalMap = new PhotonMap(maxPhotonsPerLight*bounceNum);
		causticsMap = new PhotonMap(maxCausticsPhotonsPerLight*scene.lights.size());

		this.constructMap();
		globalMap.balanceTree();
		causticsMap.balanceTree();
	}

	public void constructMap(){
		for (LightSource light : scene.lights){
			// construct the global map
			// emit photons from each light according to the light power
			Deque<Photon> photons = emitPhotons(light, false);

			int bounce = 0;
			//int frame = 1;
			while (!photons.isEmpty() && bounce < maxBounces){
				for (int i = 0; i < photons.size(); i++){
					Photon p = photons.remove();
					Ray photonRay = new Ray(p.position, p.direction.mult(-1));
					IntersectionInfo intersection = scene.getIntersection(photonRay);

					if (intersection != null){	// there is an intersection
						if (bounce == 0){
							if (//intersection.object.colourCoeffs.specularCoeffs.norm() > 0 || 
									intersection.object.colourCoeffs.transmittance.norm() > 0){
								int e1 = (int)Math.rint((1-p.direction.y)/2*255);
								int e2 = (int)Math.rint((Math.atan2(-p.direction.z,-p.direction.x) + Math.PI)
										/2/Math.PI*255);
								assert e1 >= 0 && e2 >= 0 && e1 <= 255 && e2 <= 255;

								angles[e1][e2] = true;
							}
						}
						p = p.hitSurface(intersection, globalMap, false);
						if (p != null){
							photons.add(p);
						}
					}
				}
				bounce++;

				if (bounce == 1){	// direct illumination only
					for (int p = 0; p < angles.length; p++){
						for (int t = 0; t < angles[0].length; t++){
							if (angles[p][t]){
								phiIndex.add(t);
								thetaIndex.add(p);
							}
						}
					}
				}

			}

			// construct the caustics map
			Deque<Photon> causticsPhotons = emitPhotons(light, true);

			bounce = 0;
			while (!causticsPhotons.isEmpty() && bounce < maxBounces){
				for (int i = 0; i < causticsPhotons.size(); i++){
					Photon p = causticsPhotons.remove();
					Ray photonRay = new Ray(p.position, p.direction.mult(-1));
					IntersectionInfo intersection = scene.getIntersection(photonRay);

					if (intersection != null){	// there is an intersection
						p = p.hitSurface(intersection, causticsMap, true);
						if (p != null){
							causticsPhotons.add(p);
						}
					}
				}
				bounce++;
			}
		}
	}

	public Deque<Photon> emitPhotons(LightSource light, boolean caustics){
		Deque<Photon> photons = new ArrayDeque<Photon>();

		double x, y, z;
		Vector3D direction;

		if (!caustics){
			for (int i = 0; i < maxPhotonsPerLight; i++){
				double phi = 2*Math.PI*Math.random();
				double theta = Math.acos(2*Math.random()-1);
				x = Math.cos(phi)*Math.sin(theta);
				y = -Math.cos(theta);
				z = Math.sin(phi)*Math.sin(theta);
				direction = new Vector3D(x,y,z);

				Vector3D power = light.colour.div(255);
				direction = direction.normalize();

				photons.add(new Photon(light.position, direction.mult(-1), power, true));
			}
		}
		else if (phiIndex.size() > 0){
			for (int i = 0; i < maxCausticsPhotonsPerLight; i++){
				int cell = (int)(Math.random()*phiIndex.size());
				assert phiIndex.size() == thetaIndex.size();
				double phi = (phiIndex.get(cell) + Math.random())/255*2*Math.PI - Math.PI;
				double theta = Math.acos(2*(thetaIndex.get(cell) + Math.random())/255 - 1);
				x = Math.cos(phi)*Math.sin(theta);
				y = Math.cos(theta);
				z = Math.sin(phi)*Math.sin(theta);

				direction = new Vector3D(x,y,z);
				Vector3D power = light.colour.div(255);
				direction = direction.normalize();

				photons.add(new Photon(light.position, direction.mult(-1), power, true));
			}
		}

		return photons;
	}

	public boolean renderImage(String filename, boolean flag){
		BufferedImage sceneImage = new BufferedImage(hPixelsNum,vPixelsNum,
				BufferedImage.TYPE_INT_RGB);
		int[] hitMap = computeColourMap(flag);

		// produce a RGB image from the hitMap
		sceneImage.setRGB(0, 0, hPixelsNum, vPixelsNum, hitMap, 0, hPixelsNum);
		// save the image on disk in the defined location
		try {
			File outputfile = new File(filename + ".png");
			ImageIO.write(sceneImage, "png", outputfile);
			return true;
		} catch (IOException e) {
			return false;
		}
	}

	private int[] computeColourMap(boolean mapView){		
		int[] hitMap = new int[hPixelsNum*vPixelsNum];

		// trace a ray through each pixel on the screen
		for (int row = 0; row < vPixelsNum; row++){
			for (int column = 0; column < hPixelsNum; column++){
				Vector3D direction = new Vector3D(new double[]{
						-screenWidth/2d + (2*column + 1)*pixelCentreX,
						screenHeight/2d - (2*row + 1)*pixelCentreY,
						screenFromCamera
				});
				Ray cameraRay = new Ray(camera,direction);
				// compute the powerMap
				Colour colour = trace(cameraRay,0, mapView);
				hitMap[row*hPixelsNum + column] = colour.toRGB();
			}
		}
		return hitMap;
	}

	private Colour trace(Ray ray, int bounceCount, boolean mapView){		
		// find the closest intersection
		IntersectionInfo intersection = scene.getIntersection(ray);

		// if there is an intersection
		if (intersection != null){
			// compute the shading
			Colour colour = shade(intersection, ray, bounceCount, mapView, radius);
			return colour;
		}
		else{
			if (bounceCount == 0){
				return new Colour(Color.WHITE.getRGB());
			}
			return new Colour(Color.BLACK.getRGB());
		}
	}

	private Colour shade(IntersectionInfo intersection, Ray viewer, int bounceCount, boolean mapView,
			double radius){
		ColourCoeffs objCoeffs = intersection.object.colourCoeffs;
		Vector3D normal = intersection.normal;
		Vector3D viewerFromPoint = viewer.origin.sub(intersection.position);
		viewerFromPoint = viewerFromPoint.div(Math.sqrt(viewerFromPoint.norm()));

		double red = 0, green = 0, blue = 0;		

		// find the contribution from each light source
		for (LightSource light : scene.lights){

			Colour lightColour = light.colour;
			Vector3D lightFromPoint = light.position.sub(intersection.position);
			lightFromPoint = lightFromPoint.div(Math.sqrt(lightFromPoint.norm()));

			double[] diffuseColour = diffuseShading(light, intersection, bounceCount, mapView);

			if (mapView){
				return new Colour(diffuseColour);
			}

			double[] specularColour = new double[]{0,0,0};
			if (!inShadow(intersection, light)){
				specularColour = specularShading(lightColour, objCoeffs, 
						normal, lightFromPoint, viewerFromPoint);
			}
			red += diffuseColour[Params.X] + specularColour[Params.X];
			green += diffuseColour[Params.Y] + specularColour[Params.Y];
			blue += diffuseColour[Params.Z] + specularColour[Params.Z];
		}

		// find any reflections
		if (objCoeffs.specularCoeffs.norm() > 0){	// specular surface
			if (bounceCount+1 <= maxBounces){
				Vector3D reflectionDirection = normal.mult(2*viewerFromPoint.dot(normal)).sub(viewerFromPoint);
				Ray reflectionRay = new Ray(intersection.position, reflectionDirection);
				Colour reflectedColour = trace(reflectionRay, bounceCount+1, mapView);

				double[] specularColour = specularShading(reflectedColour, objCoeffs, 
						normal, reflectionDirection, viewerFromPoint);
				red += specularColour[Params.X];
				green += specularColour[Params.Y];
				blue += specularColour[Params.Z];
			}
		}

		// find any refractions
		if (objCoeffs.transmittance.norm() > 0){	// transparent surface
			Vector3D refractedDirection = refract(viewerFromPoint.mult(-1), normal, 1/objCoeffs.refractiveIndex);
			Ray refractionRay = new Ray(intersection.position, refractedDirection);
			IntersectionInfo nextIntersection = scene.getIntersection(refractionRay);

			int internalBounces = 1;
			while(nextIntersection != null && intersection.object == nextIntersection.object && 
					internalBounces + bounceCount < maxBounces){
				// ray travels inside the object				

				refractedDirection = refract(refractedDirection, nextIntersection.normal, 
						objCoeffs.refractiveIndex);
				refractionRay = new Ray(nextIntersection.position, refractedDirection);
				nextIntersection = scene.getIntersection(refractionRay);
				internalBounces++;	

			}

			if (nextIntersection != null && nextIntersection.object != intersection.object){
				// ray comes back out of the object
				Colour refractedColour = trace(refractionRay, internalBounces + bounceCount, mapView);

				red += refractedColour.get(Params.X)*Math.pow(objCoeffs.getTransmittance(Params.X), internalBounces);
				green += refractedColour.get(Params.Y)*Math.pow(objCoeffs.getTransmittance(Params.Y), internalBounces);
				blue += refractedColour.get(Params.Z)*Math.pow(objCoeffs.getTransmittance(Params.Z), internalBounces);
			}
		}
		return new Colour(new double[]{red, green, blue});
	}

	private double[] diffuseShading(LightSource light, IntersectionInfo intersection, int bounces, boolean mapView){
		Vector3D lightFromPoint = light.position.sub(intersection.position).normalize();

		double[] primary = new double[]{0,0,0};
		if (!inShadow(intersection, light)){
			primary = super.diffuseShading(light.colour, intersection.object.colourCoeffs, intersection.normal, 
					lightFromPoint);
		}

		double[] secondary = new double[]{0,0,0};
		if (!mapView){
			Vector3D secondaryVector = new Vector3D(globalMap.getRadianceEstimates(
					intersection, false, mapView, radius)).mult(255*light.power/maxPhotonsPerLight);
			if (maxCausticsPhotonsPerLight > 0 && bounces < Params.RECURSION_LIMIT){
				Vector3D causticsVector = new Vector3D(causticsMap.getRadianceEstimates(
						intersection, true, mapView, radius/4)).mult(255*light.power/
								maxCausticsPhotonsPerLight*phiIndex.size()/angles.length/angles[0].length);
				secondaryVector = secondaryVector.plus(causticsVector);
			}
			secondary = new double[]{secondaryVector.x, secondaryVector.y, secondaryVector.z};

		}
		else{
			primary = new double[]{0,0,0};
			Vector3D secondaryVector = new Vector3D(globalMap.getRadianceEstimates(
					intersection, false, mapView, radius)).mult(255);
			if (maxCausticsPhotonsPerLight > 0){
				Vector3D causticsVector = new Vector3D(causticsMap.getRadianceEstimates(
						intersection, true, mapView, radius/4)).mult(255);
				secondaryVector = causticsVector;//secondaryVector.plus(causticsVector);
			}
			secondary = new double[]{secondaryVector.x, secondaryVector.y, secondaryVector.z};
		}

		assert primary.length == Params.VECTOR3D_DIM && secondary.length == primary.length;
		double[] result = new double[Params.VECTOR3D_DIM];
		for (int i = 0; i < primary.length; i++){
			result[i] = primary[i] + secondary[i];
		}

		return result;
	}
}
